package org.fujay.commons.langchain4j.server;

import cn.hutool.core.util.ObjectUtil;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.fujay.commons.langchain4j.param.ChatMode;
import org.fujay.commons.langchain4j.param.ChatParam;
import org.fujay.commons.langchain4j.param.Message;
import org.fujay.commons.langchain4j.rag.augmentor.EnhancedRetrieval;
import org.fujay.commons.langchain4j.rag.augmentor.aggregator.RrfContentAggregator;
import org.fujay.commons.langchain4j.rag.augmentor.injector.EnhancedInjector;
import org.fujay.commons.langchain4j.rag.augmentor.retrieval.DefaultEnhancedRetrieval;
import org.fujay.commons.langchain4j.rag.augmentor.transform.EnhancedQueryTransformerImpl;
import org.fujay.commons.langchain4j.rag.constant.Lanchain4jConstant;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.Objects;
import java.util.function.Consumer;

/**
 * 聊天服务
 * 支持普通对话和RAG检索增强对话
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ChatService {

    /**
     * 对话历史存储
     */
    private final ChatMemoryStore chatMemoryStore;

    /**
     * 线程池执行器
     */
    private final ThreadPoolTaskExecutor taskExecutor;

    private final OpenSearchClient openSearchClient;

    private final EmbeddingStore<TextSegment> embeddingStore;


    public String chat(String prompt) {

        return null;
    }

    /**
     * 执行聊天对话
     * 自动根据参数决定是否启用RAG检索增强
     *
     * @param chatParam 对话参数
     * @return 响应消息流
     */
    public Flux<Message> chat(ChatParam chatParam) {
        if (chatParam == null || StringUtils.isBlank(chatParam.getUserMessage())) {
            return Flux.error(new IllegalArgumentException("聊天参数和用户消息不能为空"));
        }

        try {
            log.info("开始处理聊天请求: userId={}, memoryId={}, chatMode={}, 消息内容={}",
                    chatParam.getUserId(), chatParam.getMemoryId(), chatParam.getChatMode(),
                    chatParam.getUserMessage());

            // 是否需要RAG检索
            boolean needRagRetrieval = ChatMode.RAG.equals(chatParam.getChatMode());

            boolean needMemory = ObjectUtil.isNotEmpty(chatParam.getMemoryId());

            // 创建AI服务构建器
            AiServices<AiAssistant> aiServices = AiServices.builder(AiAssistant.class)
                    .streamingChatModel(chatParam.getStreamingChatModel());

            // 设置记忆
            if (needMemory) {

                ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(chatParam.getMaxMemory())
                        .chatMemoryStore(chatMemoryStore)
                        .build();

                aiServices.chatMemoryProvider(chatMemoryProvider);
            }

            // 如果需要RAG检索，配置检索增强器
            if (needRagRetrieval) {
                configureRetrievalAugmenter(chatParam.getChatModel(), chatParam.getEmbeddingModel(), aiServices,chatParam.getIndexName() );
                // todo: rag检索时，如果没有传递SystemMessage，需要使用内置系统提示词
            }

            // 构建AI助手
            AiAssistant assistant = aiServices.build();

            // 初始化消息接收器
            Sinks.Many<Message> sink = Sinks.many().unicast().onBackpressureBuffer();

            TokenStream tokenStream;

            if (needMemory) {

                tokenStream = assistant.streamChatWithMemory(chatParam.getUserMessage(), chatParam.getMemoryId(), chatParam.getSystemMessage());
            } else {
                tokenStream = assistant.streamChat(chatParam.getUserMessage(), chatParam.getSystemMessage());
            }
            tokenStream.onPartialResponse(text -> {
                if (StringUtils.isNotBlank(text)) {
                    // 发送token到sink
                    sink.tryEmitNext(Message.next(chatParam.getUserId(), chatParam.getMemoryId(), text));
                }
            }).onCompleteResponse(response -> {
                // 构建使用统计
                TokenUsage tokenUsage = response.tokenUsage();
                // 发送完成消息
                Message message = Message.complete(chatParam.getUserId(), chatParam.getMemoryId(),
                        response.aiMessage().text(), tokenUsage);
                tryEmitComplete(sink, message, chatParam.getDoOnComplete());
            }).onError(error -> {
                log.error(error.getMessage(), error);
                Message message = Message.error(chatParam.getUserId(), chatParam.getMemoryId(),
                        error.getMessage());
                tryEmitComplete(sink, message, chatParam.getDoOnComplete());
            }).start();

            return sink.asFlux();
        } catch (Exception e) {
            log.error("初始化聊天服务时发生错误", e);
            return Flux.error(e);
        }
    }

    /**
     * 尝试发送完成消息
     *
     * @param sink     消息接收器
     * @param complete 完成消息
     * @param consumer 消息消费者
     */
    private void tryEmitComplete(Sinks.Many<Message> sink, Message complete, Consumer<Message> consumer) {
        sink.tryEmitNext(complete);
        sink.tryEmitComplete();
        if (Objects.nonNull(consumer)) {
            taskExecutor.submit(() -> {
                try {
                    consumer.accept(complete);
                } catch (Exception e) {
                    log.error(e.getMessage(), e);
                }
            });
        }
    }

    /**
     * 配置检索增强器
     */
    private void configureRetrievalAugmenter(ChatModel languageModel, EmbeddingModel embeddingModel, AiServices<AiAssistant> aiServices,String indexName) {

        EnhancedRetrieval.Config config = EnhancedRetrieval.Config.defaultConfig(embeddingModel, embeddingStore);
        /*检索器*/
        DefaultEnhancedRetrieval retrieval = DefaultEnhancedRetrieval.builder()
                .config(config)
                .indexName(ObjectUtil.defaultIfBlank(indexName,Lanchain4jConstant.open_search_index))
                .openSearchClient(openSearchClient)
                .build();
        /*内容注入器*/
        EnhancedInjector injector = EnhancedInjector.builder().build();
        /*内容聚合器*/
        RrfContentAggregator aggregator = new RrfContentAggregator();
        /*查询转换器*/
        EnhancedQueryTransformerImpl queryTransformer = EnhancedQueryTransformerImpl.builder()
                .ChatModel(languageModel)
                .build();

        // 创建检索增强器
        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .contentRetriever(retrieval)
                .contentInjector(injector)
                .contentAggregator(aggregator)
                .queryTransformer(queryTransformer)
                .executor(taskExecutor)
                .build();

        // 配置AI服务的检索增强器
        aiServices.retrievalAugmentor(retrievalAugmentor);
    }


}
